package track_formats

import (
	"bytes"
	"encoding/base64"
	"strconv"
	"time"

	"gitee.com/general252/ebml-go/core/codecs"
	"github.com/pion/rtp"
	"gogs.smarteye.com/servere/gobvrtp/encoder"
	"gogs.smarteye.com/servere/gobvrtp/track"
)

type H265Track struct {
	VideoTrack

	VPS, SPS, PPS bytes.Buffer
	MaxDONDiff    int

	encoder *encoder.H265Encoder
	decoder *encoder.H265Decoder
}

func NewH265Track() *H265Track {

	c := &H265Track{}

	param := c.CtlParma()
	param.SetPayloadType(96)
	param.SetCodec(track.CodecIdH265)
	param.SetClockRate(90000)
	param.SetOnReadRTPFunc(c.onRTP)

	return c
}

func (c *H265Track) UnmarshalFormat(ctx *track.FormatContext) error {
	_ = c.VideoTrack.UnmarshalFormat(ctx)

	param := c.CtlParma()

	{
		param.SetPayloadType(ctx.GetRtpMap().PayloadType)
		param.SetCodec(ctx.GetRtpMap().GetCodec())
		param.SetClockRate(ctx.GetRtpMap().ClockRate)
	}

	if len(ctx.GetFMTP().Source) > 0 {
		h265Param := codecs.NewH265Param(nil, nil, nil)
		h265Param.Load(ctx.GetFMTP().Source)

		vps, sps, pps := h265Param.GetVpsSpsPps()
		c.SetVPS(vps)
		c.SetSPS(sps)
		c.SetPPS(pps)
	}

	return nil
}

// SetVPS pps 00 00 00 01 ...
func (c *H265Track) SetVPS(vps []byte) *H265Track {
	c.VPS.Reset()
	c.VPS.Write(vps)

	return c
}

// SetSPS sps 00 00 00 01 ...
func (c *H265Track) SetSPS(sps []byte) *H265Track {
	c.SPS.Reset()
	c.SPS.Write(sps)

	return c
}

// SetPPS pps 00 00 00 01 ...
func (c *H265Track) SetPPS(pps []byte) *H265Track {
	c.PPS.Reset()
	c.PPS.Write(pps)

	return c
}

// RTPMap implements Format.
func (c *H265Track) RTPMap() string {
	return "H265/90000"
}

// FMTP implements Format.
func (c *H265Track) FMTP() map[string]string {
	f := c

	fmtp := make(map[string]string)
	if f.VPS.Len() > 0 {
		fmtp["sprop-vps"] = base64.StdEncoding.EncodeToString(f.VPS.Bytes())
	}
	if f.SPS.Len() > 0 {
		fmtp["sprop-sps"] = base64.StdEncoding.EncodeToString(f.SPS.Bytes())
	}
	if f.PPS.Len() > 0 {
		fmtp["sprop-pps"] = base64.StdEncoding.EncodeToString(f.PPS.Bytes())
	}
	if f.MaxDONDiff != 0 {
		fmtp["sprop-max-don-diff"] = strconv.FormatInt(int64(f.MaxDONDiff), 10)
	}

	if false {
		param := codecs.NewH265Param(f.VPS.Bytes(), f.SPS.Bytes(), f.PPS.Bytes())
		strFMTP := param.GetFmtpString()
		_ = strFMTP
	}

	return fmtp
}

func (c *H265Track) WritePacket(t time.Time, data []byte) error {
	param := c.CtlParma()

	if c.encoder == nil {
		// 初始化编码器
		e := &encoder.H265Encoder{
			PayloadType:           c.PayloadType(),
			SSRC:                  ToPtr(c.SSRC()),
			InitialSequenceNumber: ToPtr(param.InitialSequenceNumber),
			PayloadMaxSize:        param.PayloadMaxSize,
		}
		_ = e.Init()

		c.encoder = e
	}

	packets, err := c.encoder.Encode([][]byte{data})
	if err != nil {
		return err
	}

	timestamp, _ := c.CtlParma().RtpTimeEncode(t)

	for _, pkt := range packets {
		packet := pkt
		packet.Timestamp = timestamp
		if err = c.WriteRTP(packet); err != nil {
			return err
		}
	}

	return nil
}

func (c *H265Track) onRTP(_ track.Track, pkt *rtp.Packet) {
	packets, _ := c.Detector(pkt)

	if c.decoder == nil {
		// 初始化解码器
		decoder := &encoder.H265Decoder{}
		_ = decoder.Init()

		c.decoder = decoder
	}

	for _, packet := range packets {
		if data, err := c.decoder.Decode(packet); err == nil {
			for _, nalu := range data {
				c.onFrame(packet, append([]byte{0, 0, 0, 1}, nalu...))
			}
		}
	}
}

func init() {
	RegisterTrack(track.CodecIdH265, func() track.Track {
		return NewH265Track()
	})
}
